import struct
import csv
import numpy as np
from PIL import Image

# https://blog.51cto.com/u_16099317/12325808
def read_labels(file_path):
    with open(file_path, 'rb') as file:
        # 读取魔术数和标签数量
        magic_number = struct.unpack('>I', file.read(4))[0]
        num_labels = struct.unpack('>I', file.read(4))[0]

        print(f"Magic Number: {magic_number}")
        print(f"Number of Labels: {num_labels}")

        # 读取标签数据
        labels = []
        for _ in range(num_labels):
            label = struct.unpack('B', file.read(1))[0]
            labels.append(label)

    return labels

# 写入csv文件
def write_labels_to_csv(labels, csv_file_path):
    with open(csv_file_path, mode='w', newline='') as file:
        writer = csv.writer(file)
        writer.writerow(['Label'])  # 写入表头
        for label in labels:
            writer.writerow([label])


def read_images(file_path):
    with open(file_path, 'rb') as file:
        # 读取魔术数、图像数量、行数和列数
        magic_number = struct.unpack('>I', file.read(4))[0]
        num_images = struct.unpack('>I', file.read(4))[0]
        num_rows = struct.unpack('>I', file.read(4))[0]
        num_cols = struct.unpack('>I', file.read(4))[0]

        print(f"Magic Number: {magic_number}")
        print(f"Number of Images: {num_images}")
        print(f"Image Dimensions: {num_rows}x{num_cols}")

        # 读取图像数据
        images = []
        for _ in range(num_images):
            image = []
            for _ in range(num_rows * num_cols):
                pixel = struct.unpack('B', file.read(1))[0]
                image.append(pixel)
            images.append(image)

    return np.array(images).reshape(num_images, num_rows, num_cols)


def save_images_and_labels(images, labels, output_dir):
    for i, (image, label) in enumerate(zip(images, labels)):
        image_array = np.array(image, dtype=np.uint8)
        image_pil = Image.fromarray(image_array)
        image_pil.save(f"{output_dir}/image_{i}_label_{label}.png")


if __name__ == "__main__":
    # file_path = '../build/t10k-labels-idx1-ubyte-test'
    # csv_file_path = '../build/t10k-labels.csv'
    file_path = '../build/t10k-images-idx3-ubyte-test'
    csv_file_path = '../build/t10k-images.csv'

    labels = read_labels(file_path)
    print("First 10 labels:", labels[:10])

    write_labels_to_csv(labels, csv_file_path)
    print(f"Labels have been written to {csv_file_path}")

    # labels_file_path = 't10k-labels-idx1-ubyte'
    # images_file_path = 't10k-images-idx3-ubyte'
    # output_dir = 'mnist_images'
    #
    # labels = read_labels(labels_file_path)
    # images = read_images(images_file_path)
    #
    # save_images_and_labels(images, labels, output_dir)
    #
    # print("Images have been saved to the directory:", output_dir)

